{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 官方\n",
    "https://www.mindspore.cn/tutorials/application/zh-CN/r2.3.0rc2/cv/cnnctc.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cnnctc_dataset.tar.gz (545.4 MB)\n",
      "\n",
      "file_sizes:  60%|█████████████████▌           | 346M/572M [08:12<05:17, 713kB/s]"
     ]
    }
   ],
   "source": [
    "from download import download\n",
    "\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cnnctc_dataset.tar.gz\"\n",
    "\n",
    "download(url, \".\", kind=\"tar.gz\", replace=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore import nn\n",
    "from mindspore.common.initializer import TruncatedNormal\n",
    "\n",
    "\n",
    "class CNNCTC(nn.Cell):\n",
    "    def __init__(self, num_class, hidden_size, final_feature_width):\n",
    "        super(CNNCTC, self).__init__()\n",
    "\n",
    "        self.num_class = num_class\n",
    "        self.hidden_size = hidden_size\n",
    "        self.final_feature_width = final_feature_width\n",
    "\n",
    "        self.feature_extraction = ResNet(3, 512, BasicBlock, [1, 2, 5, 3])\n",
    "        self.prediction = nn.Dense(self.hidden_size, self.num_class)\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.feature_extraction(x)\n",
    "        # 进行维度转换，方便划分图片时按列划分\n",
    "        x = ops.transpose(x, (0, 3, 1, 2))  # [b, c, h, w] -> [b, w, c, h]\n",
    "        x = ops.reshape(x, (-1, self.hidden_size))\n",
    "        x = self.prediction(x)\n",
    "        x = ops.reshape(x, (-1, self.final_feature_width, self.num_class))\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class ResNet(nn.Cell):\n",
    "    def __init__(self, input_channel, output_channel, block, layers):\n",
    "        super(ResNet, self).__init__()\n",
    "\n",
    "        self.output_channel_block = [int(output_channel / 4), int(output_channel / 2), output_channel, output_channel]\n",
    "        self.inplanes = int(output_channel / 8)\n",
    "\n",
    "        self.conv0_1 = nn.Conv2d(input_channel, int(output_channel / 16), kernel_size=3, stride=1, padding=1,\n",
    "                                 pad_mode='pad', weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn0_1 = nn.BatchNorm2d(int(output_channel / 16), momentum=0.1)\n",
    "        self.conv0_2 = nn.Conv2d(int(output_channel / 16), self.inplanes, kernel_size=3, stride=1, padding=1,\n",
    "                                 pad_mode='pad', weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn0_2 = nn.BatchNorm2d(self.inplanes, momentum=0.1)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid')\n",
    "        self.layer1 = self._make_layer(block, self.output_channel_block[0], layers[0])\n",
    "        self.conv1 = nn.Conv2d(self.output_channel_block[0], self.output_channel_block[0], kernel_size=3, stride=1,\n",
    "                               padding=1, pad_mode='pad', weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(self.output_channel_block[0], momentum=0.1)\n",
    "\n",
    "        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid')\n",
    "        self.layer2 = self._make_layer(block, self.output_channel_block[1], layers[1])\n",
    "        self.conv2 = nn.Conv2d(self.output_channel_block[1], self.output_channel_block[1], kernel_size=3, stride=1,\n",
    "                               padding=1, pad_mode='pad', weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(self.output_channel_block[1], momentum=0.1)\n",
    "        self.pad = nn.Pad(((0, 0), (0, 0), (0, 0), (2, 2)))\n",
    "\n",
    "        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=(2, 1), pad_mode='valid')\n",
    "        self.layer3 = self._make_layer(block, self.output_channel_block[2], layers[2])\n",
    "        self.conv3 = nn.Conv2d(self.output_channel_block[2], self.output_channel_block[2], kernel_size=3, stride=1,\n",
    "                               padding=1, pad_mode='pad', weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(self.output_channel_block[2], momentum=0.1)\n",
    "\n",
    "        self.layer4 = self._make_layer(block, self.output_channel_block[3], layers[3])\n",
    "        self.conv4_1 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[3], kernel_size=2,\n",
    "                                 stride=(2, 1), padding=0, pad_mode='valid', weight_init=TruncatedNormal(0.02),\n",
    "                                 has_bias=False)\n",
    "        self.bn4_1 = nn.BatchNorm2d(self.output_channel_block[3], momentum=0.1)\n",
    "        self.conv4_2 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[3], kernel_size=2, stride=1,\n",
    "                                 padding=0, pad_mode='valid', weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn4_2 = nn.BatchNorm2d(self.output_channel_block[3], momentum=0.1)\n",
    "\n",
    "    def _make_layer(self, block, planes, blocks, stride=1):\n",
    "        downsample = None\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            downsample = nn.SequentialCell(\n",
    "                [nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, padding=0,\n",
    "                           pad_mode='same', weight_init=TruncatedNormal(0.02), has_bias=False),\n",
    "                 nn.BatchNorm2d(planes * block.expansion, momentum=0.1)]\n",
    "            )\n",
    "\n",
    "        layers = [block(self.inplanes, planes, stride, downsample)]\n",
    "        self.inplanes = planes * block.expansion\n",
    "        for _ in range(1, blocks):\n",
    "            layers.append(block(self.inplanes, planes))\n",
    "\n",
    "        return nn.SequentialCell(layers)\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.conv0_1(x)\n",
    "        x = self.bn0_1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv0_2(x)\n",
    "        x = self.bn0_2(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.maxpool1(x)\n",
    "        x = self.layer1(x)\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.maxpool2(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.maxpool3(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.bn3(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.layer4(x)\n",
    "        x = self.pad(x)\n",
    "        x = self.conv4_1(x)\n",
    "        x = self.bn4_1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv4_2(x)\n",
    "        x = self.bn4_2(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class BasicBlock(nn.Cell):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
    "        super(BasicBlock, self).__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, pad_mode='pad',\n",
    "                               weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes, momentum=0.1)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, pad_mode='pad',\n",
    "                               weight_init=TruncatedNormal(0.02), has_bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes, momentum=0.1)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.downsample = downsample\n",
    "\n",
    "    def construct(self, x):\n",
    "        residual = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            residual = self.downsample(x)\n",
    "        out = ops.add(out, residual)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 实现动态学习率\n",
    "import math\n",
    "\n",
    "def dynamic_lr(steps_per_epoch, base_lr, warmup_ratio, warmup_step, epochs):\n",
    "    total_steps = steps_per_epoch * epochs\n",
    "    warmup_steps = int(warmup_step)\n",
    "    decay_steps = total_steps - warmup_steps\n",
    "    lr = []\n",
    "    for i in range(total_steps):\n",
    "        if i < warmup_steps:\n",
    "            lr_inc = (float(base_lr) - float(base_lr * warmup_ratio)) / float(warmup_steps)\n",
    "            learning_rate = float(base_lr * warmup_ratio) + lr_inc * i\n",
    "            lr.append(learning_rate)\n",
    "        else:\n",
    "            base = float(i - warmup_steps) / float(decay_steps)\n",
    "            learning_rate = (1 + math.cos(base * math.pi)) / 2 * base_lr\n",
    "            lr.append(learning_rate)\n",
    "    return lr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import mindspore as ms\n",
    "from mindspore import ops, ParameterTuple, Tensor, load_checkpoint, load_param_into_net\n",
    "from mindspore import dtype\n",
    "from mindspore.dataset import MindDataset\n",
    "\n",
    "device_target = 'CPU'\n",
    "train_dataset_path = './cnnctc_dataset/STMJ_train_dataset.mindrecord'\n",
    "batch_size = 64\n",
    "epochs = 5\n",
    "pred_trained = ''\n",
    "final_feature_width = 26\n",
    "num_class = 37\n",
    "hidden_size = 512\n",
    "\n",
    "ms.set_context(device_target=device_target, mode=ms.PYNATIVE_MODE)\n",
    "\n",
    "# 数据集加载\n",
    "data = MindDataset(dataset_files=train_dataset_path)\n",
    "dataset_size = data.get_dataset_size()\n",
    "print('dataset_size:', dataset_size)\n",
    "print('batch_size:', batch_size)\n",
    "\n",
    "# 模型初始化\n",
    "net = CNNCTC(num_class, hidden_size, final_feature_width)\n",
    "\n",
    "# 加载预训练权重\n",
    "if pred_trained:\n",
    "    param_dict = load_checkpoint(pred_trained)\n",
    "    load_param_into_net(net, param_dict)\n",
    "    print('Pred_trained parameters loaded!')\n",
    "\n",
    "# 损失函数\n",
    "loss_fn = nn.CTCLoss(blank=0, reduction='mean', zero_infinity=False)\n",
    "\n",
    "# 学习率\n",
    "lr = Tensor(dynamic_lr(dataset_size // batch_size, 0.0005, 0.0625, 2000, epochs), dtype.float32)\n",
    "lr = lr[::-1]\n",
    "\n",
    "# 优化器\n",
    "opt = nn.RMSProp(params=net.trainable_params(),\n",
    "                 centered=True,\n",
    "                 learning_rate=lr,\n",
    "                 momentum=0.8)\n",
    "\n",
    "# 前向过程\n",
    "def forward_fn(data):\n",
    "    logits = net(data['img'])\n",
    "    logits = ops.log_softmax(logits)\n",
    "    logits = ops.transpose(logits, (1, 0, 2))\n",
    "    loss = loss_fn(logits, data['text'], data['sequence_length'], data['target_lengths'])\n",
    "    return loss, logits\n",
    "\n",
    "# 计算梯度\n",
    "grad_fn = ms.value_and_grad(forward_fn, None, ParameterTuple(net.trainable_params()), has_aux=True)\n",
    "\n",
    "# 梯度更新\n",
    "def train_step(data):\n",
    "    (loss, _), grads = grad_fn(data)\n",
    "    opt(grads)\n",
    "    return loss\n",
    "\n",
    "# 全部数据训练\n",
    "def train_loop(data_1):\n",
    "    # 将数据组织batch_size\n",
    "    batch_list = np.arange(0, dataset_size, batch_size)\n",
    "    sequence_length = np.array([final_feature_width] * batch_size, dtype=np.int32)\n",
    "    sequence_length = Tensor.from_numpy(sequence_length)\n",
    "    data_input = {}\n",
    "    for step, data in enumerate(data_1.create_dict_iterator()):\n",
    "        if step in batch_list:\n",
    "            step_b = step\n",
    "            data_img = data['img']\n",
    "            text = data['text']\n",
    "            target_lengths = data['target_lengths']\n",
    "            continue\n",
    "        if step <= step_b + batch_size - 1:\n",
    "            data_img = ops.concat((data_img, data['img']), axis=0)\n",
    "            text = ops.concat((text, data['text']), axis=0)\n",
    "            target_lengths = ops.concat((target_lengths, data['target_lengths']), axis=0)\n",
    "        if step == step_b + batch_size - 1:\n",
    "            data_input['img'] = data_img\n",
    "            data_input['text'] = text\n",
    "            data_input['target_lengths'] = target_lengths\n",
    "            data_input['sequence_length'] = sequence_length\n",
    "\n",
    "            begin_time = time.time()\n",
    "            loss = train_step(data_input)\n",
    "            end_time = time.time()\n",
    "            times = end_time - begin_time\n",
    "            if step % 1 == 0:\n",
    "                loss = loss.asnumpy()\n",
    "                print(f\"Epoch:[{int(epoch + 1):>3d}/{int(epochs):>3d}], \"\n",
    "                      f\"step:[{int(step // batch_size):>4d}/{int(dataset_size / batch_size):>4d}], \"\n",
    "                      f\"loss:{loss:>4f} , \"\n",
    "                      f\"time:{times:>3f}s, \"\n",
    "                      f\"lr:{lr[int(step // batch_size)].asnumpy():>6f}\")\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    train_loop(data)\n",
    "    ms.save_checkpoint(net, f\"./cnnctc_{epoch+1}_{dataset_size // batch_size}.ckpt\")\n",
    "print(\"Done!\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 权重下载\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cnnctc.ckpt\"\n",
    "download(url, \"cnnctc.ckpt\", replace=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 文本与数据编码\n",
    "class CTCLabelConverter():\n",
    "    def __init__(self, character):\n",
    "        dict_character = list(character)\n",
    "        self.dict = {}\n",
    "        for i, char in enumerate(dict_character):\n",
    "            self.dict[char] = i + 1\n",
    "        self.character = ['[blank]'] + dict_character\n",
    "        self.dict['[blank]'] = 0\n",
    "\n",
    "    def encode(self, text):\n",
    "        # 将文本转换为数字编码\n",
    "        length = [len(s) for s in text]\n",
    "        text = ''.join(text)\n",
    "        text = [self.dict[char] for char in text]\n",
    "\n",
    "        return np.array(text), np.array(length)\n",
    "\n",
    "    def decode(self, text_index, length):\n",
    "        # 将数字编码转换为文本\n",
    "        texts = []\n",
    "        index = 0\n",
    "        for l in length:\n",
    "            t = text_index[index:index + l]\n",
    "            char_list = []\n",
    "            for i in range(l):\n",
    "                if t[i] != self.dict['[blank]'] and (\n",
    "                        not (i > 0 and t[i - 1] == t[i])):\n",
    "                    char_list.append(self.character[t[i]])\n",
    "            text = ''.join(char_list)\n",
    "            texts.append(text)\n",
    "            index += l\n",
    "        return texts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证过程\n",
    "# 参数设置\n",
    "checkpoint_path = './cnnctc.ckpt'\n",
    "test_batch_size = 64\n",
    "character = '0123456789abcdefghijklmnopqrstuvwxyz'\n",
    "num_class = 37\n",
    "hidden_size = 512\n",
    "final_feature_width = 26\n",
    "\n",
    "# 加载测试数据集\n",
    "test_dataset_path = './cnnctc_dataset/IIIT_eval_dataset.mindrecord'\n",
    "dataset_eval = MindDataset(dataset_files=test_dataset_path)\n",
    "\n",
    "# 初始化模型\n",
    "net = CNNCTC(num_class, hidden_size, final_feature_width)\n",
    "net.set_train(False)\n",
    "\n",
    "# 加载权重文件\n",
    "param_dict = load_checkpoint(checkpoint_path)\n",
    "load_param_into_net(net, param_dict)\n",
    "print('Parameters loaded! from: ', checkpoint_path)\n",
    "\n",
    "# 初始化文本编码函数\n",
    "converter = CTCLabelConverter(character)\n",
    "\n",
    "count = 0\n",
    "correct_count = 0\n",
    "begin_time = time.time()\n",
    "num_b = 0\n",
    "\n",
    "for data in dataset_eval.create_tuple_iterator():\n",
    "    if num_b == 0 or num_b % test_batch_size == 0:\n",
    "        img, length, text = data\n",
    "        num_b = num_b + 1\n",
    "        continue\n",
    "    else:\n",
    "        img_1, length_1, text_1 = data\n",
    "        img = ops.concat((img, img_1), axis=0)\n",
    "        text = ops.concat((text, text_1), axis=0)\n",
    "        length = ops.concat((length, length_1), axis=0)\n",
    "        if num_b != 0 and (num_b + 1) % test_batch_size == 0:\n",
    "            img_tensor = Tensor(img, dtype.float32)\n",
    "            model_predict = net(img_tensor)\n",
    "\n",
    "            model_predict_1 = np.squeeze(model_predict.asnumpy())\n",
    "            preds_size = np.array([model_predict_1.shape[1]] * test_batch_size)\n",
    "            preds_index = np.argmax(model_predict_1, 2)\n",
    "            preds_index = np.reshape(preds_index, [-1])\n",
    "            preds_str = converter.decode(preds_index, preds_size)\n",
    "            label_str = converter.decode(text.asnumpy(), length.asnumpy())\n",
    "            for pred, label in zip(preds_str, label_str):\n",
    "                if pred == label:\n",
    "                    correct_count += 1\n",
    "                count += 1\n",
    "        num_b = num_b + 1\n",
    "\n",
    "times = time.time() - begin_time\n",
    "print(f'accuracy: {correct_count/count:.4f}\\n',\n",
    "      f\"time:{times:.2f}s\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    " # 推理过程\n",
    "infer_file_path = './cnnctc_dataset/predict.png'\n",
    "\n",
    "img_PIL = Image.open(infer_file_path).convert('RGB')\n",
    "img = img_PIL.resize((100, 32), resample=3)\n",
    "img = np.array(img, dtype=np.float32)\n",
    "img = np.expand_dims(img, axis=0)  # (1,32,100,3)\n",
    "img = np.transpose(img, [0, 3, 1, 2])  # (1,3,32,100)\n",
    "img_tensor = Tensor.from_numpy(img)\n",
    "\n",
    "model_predict = net(img_tensor)\n",
    "preds_size = np.array([model_predict.shape[1]])\n",
    "preds_index = np.argmax(model_predict, 2)\n",
    "preds_index = np.reshape(preds_index, [-1])\n",
    "preds_str = converter.decode(preds_index, preds_size)\n",
    "print('Predict: ', preds_str)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img_PIL)\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
